# Extracting wind intensity time series, 
# Romdhani et al 2022 (Journal of Advances in Modeling Earth Systems)


from netCDF4 import Dataset
from wrf import getvar
import numpy as np
import math
import csv
import os


#-------------------------------------------------
def create_file (Output_Dir, File_name):
	
    os.chdir(Output_Dir)
    Files= []
    test = False
    for f in os.listdir(Output_Dir):
        if f=='.DS_Store':
            pass
        else:
            Files.append(f)
    for file in Files:
        if file == File_name:
            test = True
            break
    if test == False:
        os.mkdir (File_name)
    os.chdir (Output_Dir + File_name)

#-------------------------------------------------

# This function returns a list of all wrf ouput files in the directory.
def list_ncfiles(Dir, ncfiles):
	for f in os.listdir(Dir):
		if f.startswith('wrfout'):
			ncfiles.append(f)
	ncfiles.sort()
	return (ncfiles)
#-------------------------------------------------

# Deine your input and output directories
Input_Dir = ''
Output_Dir = ''
# Choose between : 'Florence', 'Irma', 'Katrina', 'Gustav', 'Maria'
HNS = ['Katrina']
# Choose the simulation duration.
HR = ['30']
# Grid resolution
GSS = ['32km']
# Choose between : 'NoTurb', 'Smag2D', 'TKE2D'
TMS = ['NoTurb', 'Smag2D', 'TKE2D']
# Choose between : 'cLh0p25', 'cLh0p5', 'cLh1p0', 'cLh1p5'
CLHS = ['cLh1p0']
# Identify the time step
Time_Step = 6
# Identify the time index
Time_Idx = 0

# Define all hurricane's settings 
Hurricane_Settings = []

Idx = 0
for HN in HNS:
    for GS in GSS:
        for TM in TMS:
            for CLH in CLHS:

                # This list will contain all the csv_files containing the data for different times.
                csv_files = []
                Hurricane_Setting = HN + '_1Nest_' + HR[Idx] + 'Hours_MainGrid' + GS + '_' + TM + '_vert42_' + CLH + '_cfl2p0'
                print (Hurricane_Setting)

                Input_Dir_1 = Input_Dir + HN + '/' + Hurricane_Setting
                ncfiles = []
                All_Max_WND_SPD_10 = []
                ncfiles = list_ncfiles (Input_Dir_1, ncfiles)
                Times = []
                j = 0

                for ncfile in ncfiles[:]:
                    os.chdir(Input_Dir_1)
                    ncfile = Dataset(ncfile)

                    # Get U and V components of wind intensity at 10m of altitude.
                    U10_2D = np.array(getvar(ncfile, "U10", Time_Idx))
                    V10_2D = np.array(getvar(ncfile, "V10", Time_Idx))
                    U10_1D = U10_2D.flatten()
                    V10_1D = V10_2D.flatten()
                    WND_SPD_10 = U10_1D
                    # Calculate the wind intensity at each point of the map.
                    for i in range (WND_SPD_10.size - 1):
                        WND_SPD_10[i] = math.sqrt((U10_1D[i]**2)+(V10_1D[i]**2))
                    # Search for the maximum wind intensity at aspecific time step. 
                    WND_SPD_10_max = np.amax(WND_SPD_10)

                    # List the maximum wind intensity for all time steps.   
                    All_Max_WND_SPD_10.append(WND_SPD_10_max)

                    # list all the time steps
                    Times.append(Time_Step*j)
                    j = j+1

                ################################
                # Exporting the simulated data #
                ################################

                create_file (Output_Dir, HN)
                create_file (Output_Dir + HN + '/', GS)
                Path = os.getcwd()
                print (Path)

                MyFile=open('%s.csv' %Hurricane_Setting,'w')

                MyFile.write (Hurricane_Setting + "\n")
                MyFile.write ("All_Max_WND_SPD_10\n")
                for element in All_Max_WND_SPD_10:
                    MyFile.write (str(element) + "\n")
                MyFile.close()
    Idx = Idx + 1

